﻿using System.Collections.Concurrent;
using System.Reflection;
using System.Runtime.Loader;

namespace iTool.Core
{
    public class iAssemblyExtension
    {
        static ConcurrentDictionary<string, AssemblyLoadContext> AssemblyContextDictionary;

        static iAssemblyExtension()
        {
            AssemblyContextDictionary = new ConcurrentDictionary<string, AssemblyLoadContext>();
        }


        public static Assembly[] LoadDirectory(string domain, string directoryPath, string? prefix = null)
        {

            var dlls = Directory.GetFiles(directoryPath, "*.dll", SearchOption.TopDirectoryOnly);
            Assembly[] result = new Assembly[dlls.Length];

            var assemblyContext = new AssemblyLoadContext(domain, true);

            for (int i = 0; i < dlls.Length; i++)
            {
                string item = Path.GetFileName(dlls[i]);
                if (!string.IsNullOrWhiteSpace(prefix) && !item.StartsWith(prefix))
                {
                    continue;
                }
                using (var fileStream = File.OpenRead(directoryPath +"/"+ item))
                {
                    var assembly = assemblyContext.LoadFromStream(fileStream);
                    result[i] = assembly;
                }
            }

            AssemblyContextDictionary.AddOrUpdate(domain, assemblyContext, (a, b) =>
            {
                b?.Unload();
                return assemblyContext;
            });

            return result.Where(item => item != null).ToArray();
        }

        public static Assembly Load(string domain, string dllPath)
        {
            var assemblyContext = new AssemblyLoadContext(domain, true);
            var fileStream = File.OpenRead(dllPath);
            var assembly = assemblyContext.LoadFromStream(fileStream);
            fileStream.Dispose();
            AssemblyContextDictionary.AddOrUpdate(domain, assemblyContext, (a, b) =>
            {
                b?.Unload();
                return assemblyContext;
            });
            return assembly;
        }


        public static void Unload(string domain)
        {
            if (AssemblyContextDictionary.TryGetValue(domain, out AssemblyLoadContext? assemblyLoadContext))
            {
                assemblyLoadContext?.Unload();
                AssemblyContextDictionary.TryRemove(domain, out _);
            }
        }
    }
}
